from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List
import json

app = FastAPI()


# 管理连接的 类
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manage = ConnectionManager()


# 常规 HTTP路由
@app.get("/")
async def get():
    return {"message": "WebSocket 服务器 正在运行"}

# WebSocket 路由
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket,client_id:str):
    await manage.connect(websocket)
    try:
        while True:
            data=await websocket.receive_text()
            # 构造消息
            message={
                "client_id":client_id,
                "message":data
            }
            # 广播消息给所有连接
            await manage.broadcast(json.dumps(message))
    except WebSocketDisconnect:
        manage.disconnect(websocket)
        message={
            "client_id":client_id,
            "message":"离开了聊天室"
        }
        await manage.broadcast(json.dumps(message))




if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)
